# -*- coding: utf-8 -*-

import os

def file_extension(path):
  return os.path.splitext(path)[1]

if __name__ == '__main__':
    dir = r"D:\zjk\2019暑假\code\Pytorch-UNet-master\data\train"
    # txt_name = "test_pair.lst"
    res_name = "train_mask.txt"
    # f = open(os.path.join(dir, txt_name),"w")
    f_r = open(os.path.join(dir, res_name),'w')
    for file in os.listdir(dir):
        print(file)
        # print(file_extension(file))
        if file_extension(file) == '.jpg':
            # f.write(file+' '+file+'\n')
            f_r.write(file+'\n')
